#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
 @Author: Y-aong
 @FileName: task_finite_state_machines.py
 @DateTime: 2024/5/16 22:22
 @Version : 3.10
 @Description: 任务状态机操作
"""
import json
import logging
from datetime import datetime

from transitions import Machine

from apis.notice.model.alarm import Alarm
from apis.orderlines.models import TaskInstance
from apis.orderlines.schema.task_schema import TaskInstanceSchema
from orderlines.execute.running_task import RunningTask
from orderlines.utils.exceptions import StateChangeException
from orderlines.utils.orderlines_enum import TaskStates
from orderlines.utils.utils import CustomEncoder
from public.base_model import session_scope
from public.logger import logger

logging.getLogger('transitions').setLevel(logging.ERROR)
logging.getLogger('transitions').setLevel(logging.WARNING)


class TaskInstanceState:
    def __init__(self, task_instance_id: str, process_id, process_instance_id):
        self.state = None
        self.process_id = process_id
        self.process_instance_id = process_instance_id
        self.task_instance_id = task_instance_id
        self.task_states = ['SUCCESS', 'FAILURE', 'STOP', 'PENDING', 'RUNNING', 'PAUSED', 'TIMEOUT', 'RETRY', 'SKIP']
        self.table_orm = TaskInstance
        self.table_schema = TaskInstanceSchema
        self.machine = Machine(self, self.task_states, initial=self.get_init_state(), name='task')
        self.running_task = RunningTask(self.process_id, self.process_instance_id)

        self.machine.add_transition(
            trigger='start',
            source='PENDING',
            dest='RUNNING',
            conditions='can_start',
            after='on_running'
        )
        self.machine.add_transition(
            trigger='pause',
            source=['PENDING', 'RUNNING'],
            dest='PAUSED',
            conditions='can_pause',
            after='on_paused'
        )
        self.machine.add_transition(
            trigger='stop',
            source=['PENDING', 'RUNNING'],
            dest='STOP',
            after='on_stop'
        )
        self.machine.add_transition(
            trigger='recover',
            source=['PAUSED'],
            dest='RUNNING',
            conditions='can_recover',
            after='on_running'
        )
        self.machine.add_transition(
            trigger='timeout',
            source=['RUNNING'],
            dest='TIMEOUT',
            after='on_timeout'
        )
        self.machine.add_transition(
            trigger='failure',
            source=['RUNNING'],
            dest='FAILURE',
            after='on_failure'
        )
        self.machine.add_transition(
            trigger='success',
            source=['RUNNING'],
            dest='SUCCESS',
            after='on_success'
        )

        self.machine.add_transition(
            trigger='skip',
            source=['FAILURE'],
            dest='SUCCESS',
            conditions='can_skip',
            after='on_skip'
        )
        self.machine.add_transition(
            trigger='retry',
            source=['FAILURE'],
            dest='RUNNING',
            conditions='can_retry',
            after='on_retry'
        )

    def get_init_state(self):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.active == 1
            ).first()
        return self.table_schema().dump(obj).get('task_status') if obj else TaskStates.grey.value

    def can_start(self):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.task_status != TaskStates.blue.value,
                self.table_orm.active == 1
            ).first()
        if not obj:
            logger.error(f'任务状态机——无法启动:{self.task_instance_id}', task_instance_id=self.task_instance_id)
            raise StateChangeException(f'任务状态机——当前没有任务实例无法运行:{self.task_instance_id}')
        return obj is not None

    def can_pause(self):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.task_status == TaskStates.blue.value,
                self.table_orm.active == 1
            ).first()
        if not obj:
            logger.error(f'任务状态机——无法暂停:{self.task_instance_id}')
            raise StateChangeException(f'任务状态机——无法暂停:{self.task_instance_id}')
        return obj is not None

    def can_recover(self):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.task_status == TaskStates.purple.value,
                self.table_orm.active == 1
            ).first()
        if not obj:
            logger.error(f'任务状态机——无法恢复:{self.task_instance_id}')
            raise StateChangeException(f'任务状态机——无法恢复:{self.task_instance_id}')
        return obj is not None

    def can_skip(self, *args, **kwargs):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.task_status == TaskStates.red.value,
                self.table_orm.active == 1
            ).first()
        if not obj:
            logger.error(f'任务状态机——无法跳过:{self.task_instance_id}')
            raise StateChangeException(f'任务状态机——无法跳过:{self.task_instance_id},参数:{args, kwargs}')
        return obj is not None

    def can_retry(self, *args, **kwargs):
        with session_scope() as session:
            obj = session.query(self.table_orm).filter(
                self.table_orm.task_instance_id == self.task_instance_id,
                self.table_orm.task_status == TaskStates.red.value,
                self.table_orm.active == 1
            ).first()
            if not obj:
                logger.error(f'任务状态机——无法重试:{self.task_instance_id}')
                raise StateChangeException(f'任务状态机——无法重试:{self.task_instance_id},参数:{args, kwargs}')
            return obj is not None

    def on_running(self):
        task_instance_info = {'task_status': TaskStates.blue.value}

        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()
        self.running_task.send_message(self.task_instance_id)

    def on_paused(self):
        task_instance_info = {
            'task_status': TaskStates.purple.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_error_info': {'notice': '任务暂停'},
        }

        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()

        self.running_task.send_message(self.task_instance_id)

    def on_stop(self):
        task_instance_info = {
            'task_status': TaskStates.yellow.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_error_info': {'notice': '流程停止'},
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()
        self.running_task.send_message(self.task_instance_id)

    def on_success(self, result, logs=None):
        task_instance_info = {
            'task_status': TaskStates.green.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_result': json.dumps(result, cls=CustomEncoder, indent=2) if isinstance(result, dict) else result,
            'logs': logs
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()
        self.running_task.send_message(self.task_instance_id)

    def on_failure(self, error_info, logs=None):
        task_instance_info = {
            'task_status': TaskStates.red.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_result': error_info,
            'logs': logs
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()
        self.insert_into_alarm()
        self.running_task.send_message(self.task_instance_id)

    def on_skip(self, error_info, logs=None):
        task_instance_info = {
            'task_status': TaskStates.green.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_result': error_info,
            'logs': logs
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()
        self.insert_into_alarm()
        self.running_task.send_message(self.task_instance_id)

    def on_retry(self, error_info, logs=None):
        task_instance_info = {
            'task_status': TaskStates.blue.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_result': error_info,
            'logs': logs
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
            session.flush()

        self.running_task.send_message(self.task_instance_id)

    def on_timeout(self, error_info, logs):
        task_instance_info = {
            'task_status': TaskStates.orange.value,
            'end_time': datetime.now().strftime('%Y-%m-%d %H-%M-%S'),
            'task_result': error_info,
            'logs': logs
        }
        with session_scope() as session:
            session.query(TaskInstance).filter(
                TaskInstance.task_instance_id == self.task_instance_id).update(task_instance_info)
        self.insert_into_alarm()
        self.running_task.send_message(self.task_instance_id)

    def insert_into_alarm(self):
        alarm_data = dict()

        with session_scope() as session:
            obj = session.query(self.table_orm).filter(self.table_orm.task_instance_id == self.task_instance_id).first()
            task_instance_info = self.table_schema().dump(obj)
            for key, val in task_instance_info.items():
                if hasattr(Alarm, key) and key != 'id':
                    alarm_data[key] = val

            obj = Alarm(**alarm_data)
            session.add(obj)
            session.flush()

    def trigger_chance(self, trigger_methods, *args, **kwargs):
        func = getattr(self, trigger_methods)
        if args or kwargs:
            return func(*args, **kwargs)
        else:
            return func()
